"""
浏览器工具封装类
提供更便捷的浏览器操作方法，增强反爬虫能力
"""
from typing import Optional, Dict, Any, Union, List
from DrissionPage import ChromiumOptions, WebPage
import random
import time


class BrowserTool:
    """浏览器工具封装类"""
    
    # 常用的User-Agent列表
    USER_AGENTS = [
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.1.1 Safari/605.1.15",
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0",
        "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/92.0.4515.107 Safari/537.36",
        "Mozilla/5.0 (iPhone; CPU iPhone OS 14_6 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.0 Mobile/15E148 Safari/604.1",
        "Mozilla/5.0 (iPad; CPU OS 14_6 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.0 Mobile/15E148 Safari/604.1",
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36 Edg/91.0.864.59"
    ]
    
    def __init__(self, 
                 headless: bool = False, 
                 user_data_dir: str = None,
                 proxy: str = None,
                 download_path: str = None,
                 rotate_user_agent: bool = False,
                 stealth_mode: bool = True):
        """
        初始化浏览器工具
        
        参数:
            headless: 是否启用无头模式
            user_data_dir: 用户数据目录路径
            proxy: 代理服务器地址，格式如 "http://127.0.0.1:7890"
            download_path: 文件下载路径
            rotate_user_agent: 是否自动轮换User-Agent
            stealth_mode: 是否启用隐身模式，隐藏自动化特征
        """
        # 创建浏览器配置对象
        self.options = ChromiumOptions()
        
        # 设置无头模式 - 更新为新API
        if headless:
            self.options.set_argument('--headless')
        
        # 仅在Linux系统下启用
        import platform
        if platform.system() == 'Linux':
            self.options.set_argument('--no-sandbox')
    
        
        # 设置用户数据目录
        if user_data_dir:
            self.options.set_user_data_dir(user_data_dir)
        
        # 设置代理
        if proxy:
            self.options.set_proxy(proxy)
            
        # 设置下载路径
        if download_path:
            self.options.set_download_path(download_path)
            
        # 启用轮换User-Agent
        self.rotate_user_agent = rotate_user_agent
        if rotate_user_agent:
            user_agent = random.choice(self.USER_AGENTS)
            self.options.set_argument("--user-agent", user_agent)
            
        # 启用隐身模式
        if stealth_mode:
            # 禁用WebDriver特征
            self.options.set_argument("--disable-blink-features=AutomationControlled")
            # 增加随机指纹特征
            self.options.set_pref("plugins.always_open_pdf_externally", random.choice([True, False]))
            
        # 代理列表和当前代理索引
        self.proxy_list = []
        self.current_proxy_index = 0
        
        # 请求间隔控制参数
        self.min_delay = 1.0
        self.max_delay = 5.0
        
        # 指数退避算法参数
        self.backoff_factor = 2
        self.backoff_max = 60
        self.current_backoff = self.min_delay
        
        # 浏览器页面对象
        self.page = None
    
    def __del__(self):
        """
        析构方法，在对象被销毁时自动关闭浏览器
        这确保了即使用户忘记调用close()方法，浏览器也会被正确关闭
        """
        # self.close()
        pass 
    
    def create_page(self) -> WebPage:
        """
        创建并返回浏览器页面对象
        
        返回:
            WebPage对象
        """
        # 正确初始化WebPage
        self.page = WebPage(chromium_options=self.options)
        return self.page
    
    def set_proxy_list(self, proxies: List[str]) -> None:
        """
        设置代理IP池
        
        参数:
            proxies: 代理服务器地址列表，格式如 ["http://1.2.3.4:8080", "http://5.6.7.8:8080"]
        """
        self.proxy_list = proxies
    
    def rotate_proxy(self) -> None:
        """
        轮换到下一个代理IP
        """
        if not self.proxy_list:
            return
            
        self.current_proxy_index = (self.current_proxy_index + 1) % len(self.proxy_list)
        new_proxy = self.proxy_list[self.current_proxy_index]
        
        # 需要重新创建浏览器才能应用新代理
        if self.page:
            self.close()
            
        self.options.set_proxy(new_proxy)
        self.create_page()
    
    def _random_delay(self) -> None:
        """
        随机延迟，模拟人类行为
        """
        delay = random.uniform(self.min_delay, self.max_delay)
        time.sleep(delay)
    
    def _apply_backoff(self) -> None:
        """
        应用指数退避算法，遇到限制自动降低请求速率
        """
        self.current_backoff = min(self.current_backoff * self.backoff_factor, self.backoff_max)
        time.sleep(self.current_backoff)
    
    def _reset_backoff(self) -> None:
        """
        重置退避算法计数器
        """
        self.current_backoff = self.min_delay
    
    def get(self, url: str, delay: bool = True, referrer: str = None) -> None:
        """
        访问指定URL
        
        参数:
            url: 网址
            delay: 是否添加随机延迟
            referrer: 来源页面URL
        """
        if not self.page:
            self.create_page()
            
        # 设置referrer
        if referrer:
            self.page.set_headers({"Referer": referrer})
            
        # 如果启用了User-Agent轮换，则每次请求前更换
        if self.rotate_user_agent:
            user_agent = random.choice(self.USER_AGENTS)
            self.page.set_headers({"User-Agent": user_agent})
            
        # 添加随机延迟
        if delay:
            self._random_delay()
            
        # 尝试访问URL
        try:
            self.page.get(url)
            # 成功访问后重置退避
            self._reset_backoff()
        except Exception as e:
            # 请求失败，应用退避策略
            self._apply_backoff()
            raise e
    
    def find_element(self, selector: str) -> Any:
        """
        查找元素
        
        参数:
            selector: 元素选择器，可以是CSS选择器或XPath
            
        返回:
            元素对象
        """
        if not self.page:
            raise ValueError("请先调用create_page()或get()方法初始化页面")
        
        # 自动判断选择器类型
        if selector.startswith('//'):
            return self.page.ele(xpath=selector)
        else:
            return self.page.ele(selector)
    
    def find_elements(self, selector: str) -> list:
        """
        查找所有匹配的元素
        
        参数:
            selector: 元素选择器，可以是CSS选择器或XPath
            
        返回:
            元素对象列表
        """
        if not self.page:
            raise ValueError("请先调用create_page()或get()方法初始化页面")
        
        # 自动判断选择器类型
        if selector.startswith('//'):
            return self.page.eles(xpath=selector)
        else:
            return self.page.eles(selector)
    
    def click(self, selector: str, human_like: bool = True) -> None:
        """
        点击元素
        
        参数:
            selector: 元素选择器
            human_like: 是否模拟人类点击（添加延迟和鼠标移动）
        """
        element = self.find_element(selector)
        
        if human_like:
            try:
                # 先将鼠标移到元素上
                self.move_to_element(selector)
                # 添加小延迟
                time.sleep(random.uniform(0.1, 0.5))
            except Exception as e:
                # 如果鼠标移动失败，记录错误但继续执行点击
                print(f"鼠标移动失败: {e}，继续尝试点击操作")
    
        element.click()
    
        if human_like:
            # 点击后的随机延迟
            time.sleep(random.uniform(0.5, 2.0))
    
    def move_to_element(self, selector: str, smooth: bool = True) -> None:
        """
        将鼠标移动到指定元素
        
        参数:
            selector: 元素选择器
            smooth: 是否使用平滑移动
        """
        element = self.find_element(selector)
        try:
            # 使用动作链移动鼠标
            if smooth:
                self.page.actions.move_to(element)
            else:
                self.page.actions.move_to(element, duration=0)
        except Exception as e:
            print(f"鼠标移动失败: {e}")
    
    def scroll(self, direction: str = "down", pixels: int = None, smooth: bool = True) -> None:
        """
        页面滚动
        
        参数:
            direction: 滚动方向，"up"、"down"、"left" 或 "right"
            pixels: 滚动像素数，如不指定则随机
            smooth: 是否平滑滚动
        """
        if not self.page:
            raise ValueError("请先调用create_page()或get()方法初始化页面")
            
        if pixels is None:
            # 如果未指定像素数，使用随机值
            pixels = random.randint(100, 800)
            
        if direction == "down":
            y_offset = pixels
            x_offset = 0
        elif direction == "up":
            y_offset = -pixels
            x_offset = 0
        elif direction == "right":
            y_offset = 0
            x_offset = pixels
        elif direction == "left":
            y_offset = 0
            x_offset = -pixels
        else:
            raise ValueError("滚动方向必须是 'up'、'down'、'left' 或 'right'")
            
        if smooth:
            # 平滑滚动，使用JavaScript模拟
            self.page.run_js(f"""
                window.scrollBy({{
                  top: {y_offset},
                  left: {x_offset},
                  behavior: 'smooth'
                }});
            """)
            # 给滚动一些时间完成
            time.sleep(random.uniform(0.1, 0.5))
        else:
            # 快速滚动
            self.page.run_js(f"window.scrollBy({x_offset}, {y_offset});")
    
    def simulate_human_reading(self, min_time: float = 5.0, max_time: float = 15.0) -> None:
        """
        模拟人类阅读行为，包括随机滚动和停留
        
        参数:
            min_time: 最短停留时间（秒）
            max_time: 最长停留时间（秒）
        """
        if not self.page:
            raise ValueError("请先调用create_page()或get()方法初始化页面")
            
        # 获取页面高度
        height = self.page.run_js("return document.body.scrollHeight;")
        
        # 计算总停留时间
        total_time = random.uniform(min_time, max_time)
        start_time = time.time()
        
        # 已滚动距离
        scrolled = 0
        
        # 模拟阅读行为，随机滚动直到时间结束
        while time.time() - start_time < total_time:
            # 随机决定是否继续滚动
            if random.random() < 0.7:  # 70%概率滚动
                # 随机滚动距离
                scroll_amount = random.randint(100, 300)
                # 确保不会滚动超过页面底部
                if scrolled + scroll_amount < height:
                    self.scroll("down", scroll_amount)
                    scrolled += scroll_amount
                    
            # 随机停顿，模拟阅读
            time.sleep(random.uniform(0.5, 2.0))
    
    def handle_cookies(self, cookies: Dict[str, str] = None) -> Dict[str, str]:
        """
        处理Cookie
        
        参数:
            cookies: 要设置的Cookie字典
            
        返回:
            当前的Cookie字典
        """
        if not self.page:
            raise ValueError("请先调用create_page()或get()方法初始化页面")
            
        if cookies:
            # 设置Cookie
            for name, value in cookies.items():
                self.page.set_cookies({name: value})
                
        # 返回当前所有Cookie
        return self.page.get_cookies()
    
    def is_captcha_present(self) -> bool:
        """
        检测页面上是否存在验证码
        
        返回:
            如果检测到验证码则返回True，否则返回False
        """
        if not self.page:
            raise ValueError("请先调用create_page()或get()方法初始化页面")
            
        # 检查常见验证码元素
        captcha_selectors = [
            "iframe[src*='captcha']",
            "iframe[src*='recaptcha']",
            "iframe[title*='captcha']",
            ".g-recaptcha",
            "#captcha",
            "img[src*='captcha']",
            ".hcaptcha",
            "#challenge-form"  # Cloudflare挑战
        ]
        
        for selector in captcha_selectors:
            try:
                if self.page.ele(selector, timeout=0.5):
                    return True
            except:
                pass
                
        return False
    
    def input_text(self, selector: str, text: str, human_like: bool = True) -> None:
        """
        在输入框中输入文本
        
        参数:
            selector: 元素选择器
            text: 要输入的文本
            human_like: 是否模拟人类输入（间隔输入）
        """
        element = self.find_element(selector)
        
        if human_like:
            # 先点击输入框
            element.click()
            # 逐字输入，模拟人类打字
            for char in text:
                element.input(char)
                # 输入间隔
                time.sleep(random.uniform(0.05, 0.2))
        else:
            element.input(text)
    
    def get_text(self, selector: str) -> str:
        """
        获取元素文本内容
        
        参数:
            selector: 元素选择器
            
        返回:
            元素的文本内容
        """
        element = self.find_element(selector)
        return element.text
    
    def get_attribute(self, selector: str, attribute: str) -> str:
        """
        获取元素属性值
        
        参数:
            selector: 元素选择器
            attribute: 属性名称
            
        返回:
            属性值
        """
        element = self.find_element(selector)
        return element.attr(attribute)
    
    def wait_for_element(self, selector: str, timeout: float = 10) -> Any:
        """
        等待元素出现
        
        参数:
            selector: 元素选择器
            timeout: 超时时间（秒）
            
        返回:
            元素对象
        """
        if not self.page:
            raise ValueError("请先调用create_page()或get()方法初始化页面")
            
        return self.page.wait.ele_displayed(selector, timeout=timeout)
    
    def execute_js(self, script: str, *args) -> Any:
        """
        执行JavaScript代码
        
        参数:
            script: JavaScript代码
            args: 传递给JavaScript的参数
            
        返回:
            JavaScript执行结果
        """
        if not self.page:
            raise ValueError("请先调用create_page()或get()方法初始化页面")
            
        return self.page.run_js(script, *args)
    
    def bypass_cloudflare(self, max_wait: int = 30) -> bool:
        """
        尝试绕过Cloudflare保护
        
        参数:
            max_wait: 最大等待时间（秒）
            
        返回:
            是否成功绕过
        """
        if not self.page:
            raise ValueError("请先调用create_page()或get()方法初始化页面")
            
        # 检查是否存在Cloudflare挑战
        if not self.page.ele("#challenge-form", timeout=1):
            return True  # 没有Cloudflare挑战
            
        # 等待Cloudflare验证完成
        start_time = time.time()
        while time.time() - start_time < max_wait:
            # 如果挑战表单消失，说明验证通过
            if not self.page.ele("#challenge-form", timeout=0.5):
                return True
                
            # 等待一段时间再检查
            time.sleep(1)
            
        return False  # 超时未通过验证
    
    def screenshot(self, save_path: str = None) -> bytes:
        """
        截取页面截图
        
        参数:
            save_path: 保存路径，如不提供则返回图片字节数据
            
        返回:
            如果没有提供保存路径，则返回图片字节数据
        """
        if not self.page:
            raise ValueError("请先调用create_page()或get()方法初始化页面")
            
        if save_path:
            return self.page.get_screenshot(path=save_path)
        else:
            return self.page.get_screenshot()
    
    def close(self) -> None:
        """关闭浏览器"""
        if self.page:
            self.page.quit()
            self.page = None 